'''
模型的训练代码 
'''
import gc
import numpy as np
import pandas as pd
import torch

from modules import PriceModel,EncoderPriceModel,MLPPriceModel

TRAIN_PATH = './data/used_car_train_20200313.csv'
TESTB_PATH = './data/used_car_testB_20200421.csv'
CHECK_POINT = './checkpoint/'
OUTPUT_PATH = './'
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

train=pd.read_csv(TRAIN_PATH, sep=' ')
testB=pd.read_csv(TESTB_PATH, sep=' ')
train['train']=1
testB['train']=0
dataset=pd.concat([train,testB],ignore_index=True)
del train,testB
gc.collect()

# 处理特征的非数值取值
dataset.replace({'notRepairedDamage':'-'}, 2, inplace=True)
dataset['notRepairedDamage']=dataset['notRepairedDamage'].map(float).astype(np.int8)

# 时间数据类型转换
dataset['regDate']=pd.to_datetime(dataset['regDate'],format="%Y%m%d",errors='coerce')
dataset['creatDate']=pd.to_datetime(dataset['creatDate'],format="%Y%m%d",errors='coerce')
dataset.drop(labels=['offerType','seller'],axis=1,inplace=True)

dataset.drop(['SaleID','name'],axis=1,inplace=True)
# 题目要求power的范围是[0,600]
# 此处进行截断
dataset['power']=dataset['power'].map(lambda x: x if x<=600 else 600)

# 处理空值
dataset.fillna({'regDate':dataset['regDate'].median()},inplace=True)
dataset.dropna(subset=['model'],inplace=True)
dataset.fillna({'bodyType':dataset['bodyType'].median()},inplace=True)
dataset.fillna({'fuelType':dataset['fuelType'].median()},inplace=True)
dataset.fillna({'gearbox':dataset['gearbox'].median()},inplace=True)

# 获取时间特征
dataset['usedtime']=(dataset['creatDate']-dataset['regDate']).map(lambda x:x.days)
dataset.drop(['creatDate','regDate'],axis=1,inplace=True)

# 调整数据类型
for key in ['model', 'bodyType', 'fuelType','gearbox']:
    dataset[key] = dataset[key].astype(np.int64)

# 保存数据集
dataset.to_csv(OUTPUT_PATH+'dataset.csv', index=False)

# 划分数据
datasetLabel=dataset[dataset['train']==1].drop(['train'],axis=1)['price'].values
datasetTrain=dataset[dataset['train']==1].drop(['train','price'],axis=1).values
datasetTest=dataset[dataset['train']==0].drop(['train','price'],axis=1).values
del dataset
gc.collect()

# 建模
print(datasetTrain.shape)
# model = EncoderPriceModel(
#     features_num=datasetTrain.shape[1],
#     features_dim=datasetTrain.shape[-1] if 3 == len(datasetTrain.shape) else 1, 
#     features_ex=10, heads=2, deepth=50)
model = MLPPriceModel(datasetTrain.shape[1], 50)

trainer = PriceModel(model)
trainer.fit(datasetTrain,datasetLabel, epoches=20, batchsize=500, machine=DEVICE, out_dir=CHECK_POINT)

print(datasetTest.shape)
# 预测
ans = trainer.predict(datasetTest, 1500,machine=DEVICE)
testID=pd.read_csv(TESTB_PATH, sep=' ')['SaleID']
re=pd.DataFrame()
re['SaleID']=testID
re['price']=ans
re['price']=re['price'].map(lambda x: 0 if x<0 else x)
re.to_csv(OUTPUT_PATH+'ans.csv', index=False)

out = trainer.predict(datasetTrain, 1500,machine=DEVICE)
predict = pd.DataFrame()
predict['label'] = datasetLabel
predict['predict'] = out
predict.to_csv(OUTPUT_PATH+'predict.csv', index=False)

print(datasetTrain[:500].shape)
# 导出为ONNX
torch.onnx.export(trainer.model,
    torch.tensor(datasetTrain[:500]).unsqueeze(-1).float(),
    f'{CHECK_POINT}{trainer.model._get_name()}.onnx',
    training=torch.onnx.TrainingMode.EVAL,
    do_constant_folding=False,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input':{0:'batch'},
        'output':{0:'batch'}
    }
)
